from ultralytics import YOLO
import os
from onnx import numpy_helper
import onnx

def export_yolov8_onnx(model_path, output_dir='./models'):
    """优化后的ONNX导出函数，确保路径正确"""
    try:
        # 创建输出目录（包含权限检查）
        os.makedirs(output_dir, exist_ok=True)
        
        # 加载模型并验证路径有效性
        model = YOLO(model_path)
        
        # 构建完整输出路径
        base_name = os.path.basename(model_path).rsplit('.', 1)[0]
        onnx_path = os.path.join(output_dir, f"{base_name}.onnx")
        
        # 显式指定导出路径（关键修复）
        model.export(
            format='onnx', 
            imgsz=640,
            simplify=True,
            batch=1,
            opset=12,
        )
        
        return onnx_path
    
    except Exception as e:
        print(f"导出失败: {str(e)}")
        return None

def convert_onnx_int64_to_int32(onnx_path):
    model = onnx.load(onnx_path)
    
    # 转换所有INT64的initializer到INT32
    for tensor in model.graph.initializer:
        if tensor.data_type == onnx.TensorProto.INT64:
            # 转换数据类型
            tensor.data_type = onnx.TensorProto.INT32
            # 转换数据内容
            if tensor.HasField("raw_data"):
                array = numpy_helper.to_array(tensor)
                array = array.astype(np.int32)
                tensor.raw_data = array.tobytes()
    
    # 转换所有INT64的输入/输出/中间变量
    for node in model.graph.node:
        for attr in node.attribute:
            if attr.type == onnx.AttributeProto.INTS and attr.ints:
                if any(isinstance(i, int) and i > 2**32 for i in attr.ints):
                    attr.ints[:] = [int(i) for i in attr.ints]
    
    # 保存转换后的模型
    new_onnx_path = onnx_path.replace(".onnx", "_int32.onnx")
    onnx.save(model, new_onnx_path)
    return new_onnx_path

if __name__ == "__main__":
    pt_model = "/home/lidar/Foursteps/yolov8/runs/train/exp9/weights/best.pt" 
    onnx_path = export_yolov8_onnx(pt_model)
    if onnx_path:
        converted_path = convert_onnx_int64_to_int32(onnx_path)
        print(f"转换后的ONNX路径: {converted_path}")
